import os
import torch.utils.data as torch_data
import torchvision
from torchvision.transforms import transforms

from tasks import Task


class FacescrubTask(Task):
    normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])
    
    _train_dst_path:str
    _test_dst_path:str

    def load_data(self):
        transform_train = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            self.normalize
        ])

        transform_test = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            self.normalize
        ])

        self.train_dataset = torchvision.datasets.ImageFolder(
            root=self._train_dst_path,
            transform=transform_train)
        

        self.train_dataloader = torch_data.DataLoader(self.train_dataset,
                                                  batch_size=self.params.BS,
                                                  shuffle=True,
                                                  num_workers=0)
        
        self.test_dataset = torchvision.datasets.ImageFolder(
            root=self._test_dst_path,
            transform=transform_test)
                
        self.test_dataloader = torch_data.DataLoader(self.test_dataset,
                                                 batch_size=self.params.BS,
                                                 shuffle=False,
                                                 num_workers=0)
        

    #override
    def _init_task(self):
        self._train_dst_path,self._test_dst_path = set_dst_path(self.params.paths["data_path"])
        mk_dst_file(self._train_dst_path,self._test_dst_path)
        super()._init_task()


def set_dst_path(data_path:str):
    """
    description:
        set dataset path
    """
    _dst_path:str = "facescrub"
    _dst_path = os.path.join(data_path,_dst_path)
    _train_dst_path = os.path.join(_dst_path,"train")
    _test_dst_path = os.path.join(_dst_path,"test")
    return _train_dst_path,_test_dst_path
        

def mk_dst_file(train_dst_path:str,test_dst_path:str):
    """
    description:
        get dataset file
    """
    if train_dst_path == None or test_dst_path == None:
        assert False, "can not find the FACESCRUB's dataset path"

    if not os.path.exists(train_dst_path):
        os.mkdir(train_dst_path)
        
    if not os.path.exists(test_dst_path):
        os.mkdir(test_dst_path)